#include "bw_tree.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>

// =============================================================================
// Node Consolidation Implementation
// =============================================================================

// Helper structure for consolidation
typedef struct consolidation_entry {
    bw_key_t* key;
    bw_value_t* value;
    bool is_deleted;
    struct consolidation_entry* next;
} consolidation_entry_t;

// Apply all deltas to create a consolidated view
static consolidation_entry_t* apply_deltas_to_entries(bw_node_t* node_chain, 
                                                     memory_manager_t* memory_mgr) {
    if (!node_chain || !memory_mgr) return NULL;
    
    consolidation_entry_t* entries = NULL;
    consolidation_entry_t* entries_tail = NULL;
    
    // First, extract base node entries
    bw_node_t* current = node_chain;
    while (current && current->type >= BW_NODE_DELTA_INSERT) {
        current = atomic_load(&current->next);
    }
    
    if (!current || current->type != BW_NODE_LEAF) {
        return NULL; // Only handle leaf nodes for now
    }
    
    // Extract entries from base leaf node
    char* data_ptr = current->data;
    size_t offset = 0;
    
    for (int i = 0; i < current->count; i++) {
        bw_kv_pair_t* pair = (bw_kv_pair_t*)(data_ptr + offset);
        
        consolidation_entry_t* entry = mm_alloc(memory_mgr, sizeof(consolidation_entry_t), 
                                               MEM_TYPE_INDEX_BUFFER);
        if (!entry) break;
        
        entry->key = bw_key_create(memory_mgr, pair->data, pair->key_size);
        entry->value = bw_value_create(memory_mgr, pair->data + pair->key_size, pair->value_size);
        entry->is_deleted = false;
        entry->next = NULL;
        
        if (!entry->key || !entry->value) {
            if (entry->key) bw_key_destroy(entry->key, memory_mgr);
            if (entry->value) bw_value_destroy(entry->value, memory_mgr);
            mm_free(memory_mgr, entry);
            break;
        }
        
        if (!entries) {
            entries = entry;
            entries_tail = entry;
        } else {
            entries_tail->next = entry;
            entries_tail = entry;
        }
        
        offset += sizeof(bw_kv_pair_t) + pair->key_size + pair->value_size;
    }
    
    // Apply deltas in reverse order (oldest to newest)
    // Build delta array first
    delta_record_t* deltas[BW_TREE_MAX_DELTA_CHAIN * 2];
    int delta_count = 0;
    
    current = node_chain;
    while (current && current->type >= BW_NODE_DELTA_INSERT && 
           delta_count < BW_TREE_MAX_DELTA_CHAIN * 2) {
        deltas[delta_count++] = (delta_record_t*)current;
        current = atomic_load(&current->next);
    }
    
    // Apply deltas from oldest to newest
    for (int i = delta_count - 1; i >= 0; i--) {
        delta_record_t* delta = deltas[i];
        
        switch (delta->type) {
            case BW_NODE_DELTA_INSERT: {
                // Find if key already exists
                consolidation_entry_t* existing = entries;
                bool found = false;
                
                while (existing) {
                    if (bw_key_compare(existing->key, delta->data.insert.key) == 0) {
                        // Update existing entry
                        if (existing->value) {
                            bw_value_destroy(existing->value, memory_mgr);
                        }
                        existing->value = bw_value_create(memory_mgr, 
                                                        delta->data.insert.value->data,
                                                        delta->data.insert.value->size);
                        existing->is_deleted = false;
                        found = true;
                        break;
                    }
                    existing = existing->next;
                }
                
                if (!found) {
                    // Add new entry
                    consolidation_entry_t* new_entry = mm_alloc(memory_mgr, 
                                                              sizeof(consolidation_entry_t),
                                                              MEM_TYPE_INDEX_BUFFER);
                    if (new_entry) {
                        new_entry->key = bw_key_create(memory_mgr,
                                                     delta->data.insert.key->data,
                                                     delta->data.insert.key->size);
                        new_entry->value = bw_value_create(memory_mgr,
                                                         delta->data.insert.value->data,
                                                         delta->data.insert.value->size);
                        new_entry->is_deleted = false;
                        new_entry->next = entries;
                        entries = new_entry;
                    }
                }
                break;
            }
            
            case BW_NODE_DELTA_DELETE: {
                // Mark entry as deleted
                consolidation_entry_t* existing = entries;
                while (existing) {
                    if (bw_key_compare(existing->key, delta->data.delete.key) == 0) {
                        existing->is_deleted = true;
                        break;
                    }
                    existing = existing->next;
                }
                break;
            }
            
            case BW_NODE_DELTA_UPDATE: {
                // Update existing entry
                consolidation_entry_t* existing = entries;
                while (existing) {
                    if (bw_key_compare(existing->key, delta->data.update.key) == 0) {
                        if (existing->value) {
                            bw_value_destroy(existing->value, memory_mgr);
                        }
                        existing->value = bw_value_create(memory_mgr,
                                                        delta->data.update.new_value->data,
                                                        delta->data.update.new_value->size);
                        existing->is_deleted = false;
                        break;
                    }
                    existing = existing->next;
                }
                break;
            }
            
            default:
                // Handle other delta types (split, merge, etc.) in full implementation
                break;
        }
    }
    
    return entries;
}

// Create new consolidated node from entries
static bw_node_t* create_consolidated_node(consolidation_entry_t* entries,
                                          memory_manager_t* memory_mgr,
                                          uint32_t node_size) {
    if (!entries || !memory_mgr) return NULL;
    
    // Count valid (non-deleted) entries and calculate size
    int valid_count = 0;
    size_t total_size = sizeof(bw_node_t);
    
    consolidation_entry_t* current = entries;
    while (current) {
        if (!current->is_deleted && current->key && current->value) {
            valid_count++;
            total_size += sizeof(bw_kv_pair_t) + current->key->size + current->value->size;
        }
        current = current->next;
    }
    
    if (total_size > node_size) {
        return NULL; // Node would be too large
    }
    
    // Create new leaf node
    bw_node_t* new_node = bw_node_create_leaf(memory_mgr, node_size);
    if (!new_node) return NULL;
    
    // Sort entries by key (simple insertion sort for small datasets)
    consolidation_entry_t** sorted_entries = mm_alloc(memory_mgr, 
                                                     valid_count * sizeof(consolidation_entry_t*),
                                                     MEM_TYPE_TEMP_BUFFER);
    if (!sorted_entries) {
        bw_node_destroy(new_node, memory_mgr);
        return NULL;
    }
    
    int sorted_count = 0;
    current = entries;
    while (current) {
        if (!current->is_deleted && current->key && current->value) {
            // Insert in sorted order
            int insert_pos = sorted_count;
            for (int i = 0; i < sorted_count; i++) {
                if (bw_key_compare(current->key, sorted_entries[i]->key) < 0) {
                    insert_pos = i;
                    break;
                }
            }
            
            // Shift entries to make room
            for (int i = sorted_count; i > insert_pos; i--) {
                sorted_entries[i] = sorted_entries[i - 1];
            }
            
            sorted_entries[insert_pos] = current;
            sorted_count++;
        }
        current = current->next;
    }
    
    // Pack entries into node data
    char* data_ptr = new_node->data;
    size_t offset = 0;
    
    for (int i = 0; i < sorted_count; i++) {
        consolidation_entry_t* entry = sorted_entries[i];
        
        bw_kv_pair_t* pair = (bw_kv_pair_t*)(data_ptr + offset);
        pair->key_size = entry->key->size;
        pair->value_size = entry->value->size;
        
        offset += sizeof(bw_kv_pair_t);
        
        // Copy key
        memcpy(data_ptr + offset, entry->key->data, entry->key->size);
        offset += entry->key->size;
        
        // Copy value
        memcpy(data_ptr + offset, entry->value->data, entry->value->size);
        offset += entry->value->size;
    }
    
    new_node->count = sorted_count;
    
    mm_free(memory_mgr, sorted_entries);
    return new_node;
}

// Clean up consolidation entries
static void cleanup_consolidation_entries(consolidation_entry_t* entries,
                                         memory_manager_t* memory_mgr) {
    while (entries) {
        consolidation_entry_t* next = entries->next;
        
        if (entries->key) bw_key_destroy(entries->key, memory_mgr);
        if (entries->value) bw_value_destroy(entries->value, memory_mgr);
        mm_free(memory_mgr, entries);
        
        entries = next;
    }
}

bw_result_t bw_node_consolidate(bw_tree_t* tree, page_id_t page_id) {
    if (!tree || page_id == 0 || atomic_load(&tree->is_destroyed)) {
        return BW_ERROR_INVALID_PARAM;
    }
    
    bw_page_t* page = mapping_table_get_page(tree->mapping_table, page_id);
    if (!page) {
        return BW_ERROR_INVALID_PARAM;
    }
    
    // Lock page for consolidation
    if (pthread_mutex_trylock(&page->mutex) != 0) {
        // Another thread is already consolidating
        atomic_fetch_sub(&page->ref_count, 1);
        return BW_SUCCESS; // Not an error, just skip
    }
    
    bw_result_t result = BW_SUCCESS;
    
    // Get current node chain
    bw_node_t* current_chain = atomic_load(&page->node_ptr);
    if (!current_chain) {
        pthread_mutex_unlock(&page->mutex);
        atomic_fetch_sub(&page->ref_count, 1);
        return BW_ERROR_INVALID_PARAM;
    }
    
    // Check if consolidation is still needed
    if (!bw_node_needs_consolidation(current_chain, tree->max_delta_chain)) {
        pthread_mutex_unlock(&page->mutex);
        atomic_fetch_sub(&page->ref_count, 1);
        return BW_SUCCESS;
    }
    
    // Apply all deltas to create consolidated entries
    consolidation_entry_t* entries = apply_deltas_to_entries(current_chain, tree->memory_mgr);
    if (!entries) {
        pthread_mutex_unlock(&page->mutex);
        atomic_fetch_sub(&page->ref_count, 1);
        return BW_ERROR_OUT_OF_MEMORY;
    }
    
    // Create new consolidated node
    bw_node_t* new_node = create_consolidated_node(entries, tree->memory_mgr, tree->node_size);
    if (!new_node) {
        cleanup_consolidation_entries(entries, tree->memory_mgr);
        pthread_mutex_unlock(&page->mutex);
        atomic_fetch_sub(&page->ref_count, 1);
        return BW_ERROR_OUT_OF_MEMORY;
    }
    
    // Install new node using CAS
    if (atomic_compare_exchange_strong(&page->node_ptr, &current_chain, new_node)) {
        // Success - defer deletion of old chain
        gc_defer_delete(tree->gc_context, current_chain);
        atomic_fetch_add(&tree->consolidations, 1);
    } else {
        // CAS failed - another thread modified the node
        bw_node_destroy(new_node, tree->memory_mgr);
        result = BW_ERROR_CONCURRENT_UPDATE;
    }
    
    cleanup_consolidation_entries(entries, tree->memory_mgr);
    pthread_mutex_unlock(&page->mutex);
    atomic_fetch_sub(&page->ref_count, 1);
    
    return result;
}

bw_node_t* bw_node_apply_deltas(bw_node_t* base_node, memory_manager_t* memory_mgr) {
    if (!base_node || !memory_mgr) return NULL;
    
    // This is a simplified version that just returns the base node
    // Full implementation would apply all deltas to create a logical view
    return base_node;
}

// =============================================================================
// Node Splitting and Merging (Simplified Implementation)
// =============================================================================

bw_result_t bw_node_split(bw_tree_t* tree, page_id_t page_id, bw_key_t* split_key) {
    if (!tree || page_id == 0 || !split_key || atomic_load(&tree->is_destroyed)) {
        return BW_ERROR_INVALID_PARAM;
    }
    
    // This is a placeholder for node splitting
    // Full implementation would:
    // 1. Lock the page
    // 2. Consolidate if needed
    // 3. Split the node into two parts
    // 4. Create new page for right half
    // 5. Update parent node with split key
    // 6. Install split delta records
    
    atomic_fetch_add(&tree->splits, 1);
    return BW_SUCCESS;
}

bw_result_t bw_node_merge(bw_tree_t* tree, page_id_t left_page_id, page_id_t right_page_id) {
    if (!tree || left_page_id == 0 || right_page_id == 0 || atomic_load(&tree->is_destroyed)) {
        return BW_ERROR_INVALID_PARAM;
    }
    
    // This is a placeholder for node merging
    // Full implementation would:
    // 1. Lock both pages
    // 2. Consolidate if needed
    // 3. Merge right node into left node
    // 4. Update parent node to remove right page reference
    // 5. Install merge delta records
    
    atomic_fetch_add(&tree->merges, 1);
    return BW_SUCCESS;
}

// =============================================================================
// Node Inspection and Debugging
// =============================================================================

void bw_node_print(bw_node_t* node, int level) {
    if (!node) {
        printf("%*sNULL node\n", level * 2, "");
        return;
    }
    
    const char* type_names[] = {
        "LEAF", "INTERNAL", "DELTA_INSERT", "DELTA_DELETE", 
        "DELTA_UPDATE", "DELTA_SPLIT", "DELTA_MERGE", "DELTA_REMOVE", "INVALID"
    };
    
    int type_idx = (node->type >= 0 && node->type < BW_NODE_INVALID) ? node->type : BW_NODE_INVALID;
    
    printf("%*sNode: type=%s, size=%u, count=%u, level=%u, epoch=%lu\n",
           level * 2, "", type_names[type_idx], node->size, node->count, 
           node->level, node->create_epoch);
    
    if (node->type == BW_NODE_LEAF) {
        // Print leaf node entries
        char* data_ptr = node->data;
        size_t offset = 0;
        
        for (int i = 0; i < node->count && i < 5; i++) { // Limit output
            bw_kv_pair_t* pair = (bw_kv_pair_t*)(data_ptr + offset);
            printf("%*s  Entry %d: key_size=%u, value_size=%u\n",
                   level * 2, "", i, pair->key_size, pair->value_size);
            offset += sizeof(bw_kv_pair_t) + pair->key_size + pair->value_size;
        }
        
        if (node->count > 5) {
            printf("%*s  ... and %u more entries\n", level * 2, "", node->count - 5);
        }
    }
    
    // Print next delta in chain
    bw_node_t* next = atomic_load(&node->next);
    if (next) {
        printf("%*sNext delta:\n", level * 2, "");
        bw_node_print(next, level + 1);
    }
}

size_t bw_node_get_key_count(bw_node_t* node) {
    if (!node) return 0;
    
    size_t count = 0;
    
    // Count keys in base node
    bw_node_t* current = node;
    while (current && current->type >= BW_NODE_DELTA_INSERT) {
        if (current->type == BW_NODE_DELTA_INSERT) {
            count++; // Insert adds a key
        } else if (current->type == BW_NODE_DELTA_DELETE) {
            // Delete removes a key (but we can't be sure it exists)
        }
        current = atomic_load(&current->next);
    }
    
    if (current && (current->type == BW_NODE_LEAF || current->type == BW_NODE_INTERNAL)) {
        count += current->count;
    }
    
    return count;
}

size_t bw_delta_chain_length(bw_node_t* node) {
    if (!node) return 0;
    
    size_t length = 0;
    bw_node_t* current = node;
    
    while (current && current->type >= BW_NODE_DELTA_INSERT) {
        length++;
        current = atomic_load(&current->next);
    }
    
    return length;
}

bool bw_node_is_consistent(bw_node_t* node) {
    if (!node) return false;
    
    // Basic consistency checks
    if (node->size < sizeof(bw_node_t)) return false;
    if (node->type >= BW_NODE_INVALID) return false;
    
    // Check delta chain length
    if (bw_delta_chain_length(node) > BW_TREE_MAX_DELTA_CHAIN * 2) {
        return false;
    }
    
    return true;
}